
import torch,os,pdb
import warnings
warnings.filterwarnings('ignore')
from diffusers import FluxFillPipeline
# from diffusers.utils import load_image
# from itertools import product
# from image_gen_aux import DepthPreprocessor

# from util_flux import pad_image
# from util_flux import horizontal_concat_images

# from util_sam import get_mask_by_sam
# from util_mask import get_erosed_mask_by_radtio,add_random_holes


FLUX_FILL='/home/shengjie/ckp/FLUX.1-Fill-dev'
FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_DEPTH_LORA='/home/shengjie/ckp/FLUX.1-Depth-dev-lora'
FLUX='/data/models/FLUX___1-dev'

DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'


target_shape = (1024,1024)

# examples_dir = '/data/shengjie/style_zhenzhi/'

pipe = FluxFillPipeline.from_pretrained(FLUX_FILL, 
                                        torch_dtype=torch.bfloat16).to("cuda")



def test():
    # torch.Size([1, 512, 4096])
    # torch.Size([1, 768])
    # torch.Size([512, 3])
    with torch.no_grad():
        prompt_embeds, pooled_prompt_embeds, text_ids  = \
                pipe.encode_prompt(['sweater'],prompt_2=None)

    # pdb.set_trace()

    torch.save(prompt_embeds,'./prompt_embeds.pth')
    torch.save(pooled_prompt_embeds,'./pooled_prompt_embeds.pth')
    torch.save(text_ids,'./text_ids.pth')
def generate_emb_from_caption(caption):
    with torch.no_grad():
        prompt_embeds, pooled_prompt_embeds, text_ids  = \
                pipe.encode_prompt(prompt=[caption],prompt_2=[caption])
    return prompt_embeds, pooled_prompt_embeds, text_ids 
def save_emb(prompt_embeds, pooled_prompt_embeds, text_ids ,
             emb_path):
    try:
        torch.save({
            'prompt_embeds':prompt_embeds,
            'pooled_prompt_embeds':pooled_prompt_embeds,
            'text_ids':text_ids,
        },emb_path)
    except:
        return False
    return True

if __name__=='__main__':
    prompt = f"The pair of images highlights first clothing showing second clothing's texture with third clothing's shape, high resolution, 4K, 8K; " \
            f"[IMAGE1] Synthesis clothing with second's texture and third's shape." \
            f"[IMAGE2] Detailed texture shot of a clothing." \
            f"[IMAGE3] Detailed shape shot of a clothing."
    prompt_embeds, pooled_prompt_embeds, text_ids = \
        generate_emb_from_caption(prompt)
    pdb.set_trace()
    import json
    with open('./zhenzhi_data.json') as f:
        data = json.load(f)
    prompt_embeds_path = data['caption']['prompt_embeds']
    pooled_prompt_embeds_path = data['caption']['pooled_prompt_embeds']
    text_ids_path = data['caption']['text_ids']
    torch.save(prompt_embeds,prompt_embeds_path)
    torch.save(pooled_prompt_embeds,pooled_prompt_embeds_path)
    torch.save(text_ids,text_ids_path)